"""
Test suite for WRF NetCDF background field reader

This test suite validates the WRFReader module functionality including:
- Reading WRF NetCDF files
- Unstaggering wind components
- Converting perturbation to total fields
- Grid metadata extraction
- Coordinate transformations
"""

using Test
using NCDatasets
using Dates

# Add parent directory to load path
push!(LOAD_PATH, joinpath(@__DIR__, "..", "src"))

using GSICoreAnalysis
using GSICoreAnalysis.DataIO
using GSICoreAnalysis.DataIO.WRFReader
using GSICoreAnalysis.GridOperations

# Test file paths
const WRF_TEST_FILE = "/home/docker/comgsi/tutorial/case_data/2018081212/bkg/wrfinput_d01.mem0001"
const WRF_AVAILABLE = isfile(WRF_TEST_FILE)

@testset "WRFReader Tests" begin

    @testset "Unstaggering Functions" begin
        # Test U unstaggering
        @testset "U unstaggering" begin
            # Create test data: 5x4x3 staggered in x-direction
            U_stag = randn(5, 4, 3)
            U_unstag = WRFReader.unstagger_u(U_stag)

            @test size(U_unstag) == (4, 4, 3)

            # Check that unstaggered value is average of adjacent staggered values
            for k in 1:3, j in 1:4, i in 1:4
                expected = 0.5 * (U_stag[i,j,k] + U_stag[i+1,j,k])
                @test U_unstag[i,j,k] ≈ expected
            end
        end

        # Test V unstaggering
        @testset "V unstaggering" begin
            # Create test data: 4x5x3 staggered in y-direction
            V_stag = randn(4, 5, 3)
            V_unstag = WRFReader.unstagger_v(V_stag)

            @test size(V_unstag) == (4, 4, 3)

            # Check that unstaggered value is average of adjacent staggered values
            for k in 1:3, j in 1:4, i in 1:4
                expected = 0.5 * (V_stag[i,j,k] + V_stag[i,j+1,k])
                @test V_unstag[i,j,k] ≈ expected
            end
        end

        # Test vertical unstaggering
        @testset "Vertical unstaggering" begin
            # Create test data: 4x4x5 staggered in z-direction
            field_stag = randn(4, 4, 5)
            field_unstag = WRFReader.unstagger_vertical(field_stag)

            @test size(field_unstag) == (4, 4, 4)

            # Check that unstaggered value is average of adjacent staggered values
            for k in 1:4, j in 1:4, i in 1:4
                expected = 0.5 * (field_stag[i,j,k] + field_stag[i,j,k+1])
                @test field_unstag[i,j,k] ≈ expected
            end
        end
    end

    @testset "WRF Time Parsing" begin
        # Test standard WRF time format
        time_str = "2018-08-12_12:00:00"
        dt = WRFReader.parse_wrf_time(time_str)

        @test dt == DateTime(2018, 8, 12, 12, 0, 0)

        # Test with whitespace
        time_str_ws = "  2021-01-15_06:30:45  "
        dt_ws = WRFReader.parse_wrf_time(time_str_ws)

        @test dt_ws == DateTime(2021, 1, 15, 6, 30, 45)
    end

    @testset "WRF Projection Mapping" begin
        @test WRFReader.get_wrf_projection_name(1) == "lambert_conformal"
        @test WRFReader.get_wrf_projection_name(2) == "polar_stereographic"
        @test WRFReader.get_wrf_projection_name(3) == "mercator"
        @test WRFReader.get_wrf_projection_name(6) == "latlon"

        # Unknown projection should default to lambert_conformal with warning
        @test WRFReader.get_wrf_projection_name(99) == "lambert_conformal"
    end

    @testset "Mock WRF File Reading" begin
        # Create a mock WRF-like NetCDF file for testing
        mock_file = joinpath(@__DIR__, "mock_wrf.nc")

        # Create mock WRF file
        NCDataset(mock_file, "c") do ds
            # Define dimensions
            defDim(ds, "Time", 1)
            defDim(ds, "west_east", 10)
            defDim(ds, "south_north", 8)
            defDim(ds, "bottom_top", 5)
            defDim(ds, "west_east_stag", 11)
            defDim(ds, "south_north_stag", 9)
            defDim(ds, "bottom_top_stag", 6)
            defDim(ds, "DateStrLen", 19)

            # Create coordinate variables
            defVar(ds, "XLAT", Float32, ("west_east", "south_north", "Time"))
            defVar(ds, "XLONG", Float32, ("west_east", "south_north", "Time"))
            defVar(ds, "Times", Char, ("DateStrLen", "Time"))

            # Create atmospheric variables
            defVar(ds, "U", Float32, ("west_east_stag", "south_north", "bottom_top", "Time"))
            defVar(ds, "V", Float32, ("west_east", "south_north_stag", "bottom_top", "Time"))
            defVar(ds, "T", Float32, ("west_east", "south_north", "bottom_top", "Time"))
            defVar(ds, "QVAPOR", Float32, ("west_east", "south_north", "bottom_top", "Time"))
            defVar(ds, "P", Float32, ("west_east", "south_north", "bottom_top", "Time"))
            defVar(ds, "PB", Float32, ("west_east", "south_north", "bottom_top", "Time"))
            defVar(ds, "PSFC", Float32, ("west_east", "south_north", "Time"))
            defVar(ds, "PH", Float32, ("west_east", "south_north", "bottom_top_stag", "Time"))
            defVar(ds, "PHB", Float32, ("west_east", "south_north", "bottom_top_stag", "Time"))

            # Set global attributes
            ds.attrib["MAP_PROJ"] = 1
            ds.attrib["TRUELAT1"] = 30.0
            ds.attrib["TRUELAT2"] = 60.0
            ds.attrib["STAND_LON"] = -95.0
            ds.attrib["CEN_LAT"] = 40.0
            ds.attrib["CEN_LON"] = -95.0
            ds.attrib["DX"] = 9000.0
            ds.attrib["DY"] = 9000.0
            ds.attrib["TITLE"] = "Mock WRF V4.3 MODEL"

            # Fill with mock data
            nx, ny, nz = 10, 8, 5

            # Coordinates
            lat_data = Float32.(repeat(30.0:0.5:33.5, 1, nx)')
            lon_data = Float32.(repeat(-100.0:0.5:-95.5, ny, 1))
            ds["XLAT"][:, :, 1] = lon_data
            ds["XLONG"][:, :, 1] = lat_data
            ds["Times"][:, 1] = collect("2018-08-12_12:00:00")

            # Winds (realistic values around 10 m/s)
            ds["U"][:, :, :, 1] = Float32.(5.0 .+ 5.0 * randn(11, ny, nz))
            ds["V"][:, :, :, 1] = Float32.(3.0 .+ 4.0 * randn(nx, 9, nz))

            # Temperature perturbation (around 0-10K from base state)
            ds["T"][:, :, :, 1] = Float32.(5.0 .+ 3.0 * randn(nx, ny, nz))

            # Water vapor (realistic values 0-0.01 kg/kg)
            ds["QVAPOR"][:, :, :, 1] = Float32.(abs.(0.005 .+ 0.003 * randn(nx, ny, nz)))

            # Pressure (base + perturbation, around 100000 Pa at surface)
            # Decreasing with height
            for k in 1:nz
                pressure_level = 100000.0 - (k-1) * 15000.0
                ds["PB"][:, :, k, 1] .= Float32(pressure_level)
                ds["P"][:, :, k, 1] = Float32.(100.0 .+ 50.0 * randn(nx, ny))
            end

            # Surface pressure
            ds["PSFC"][:, :, 1] = Float32.(101325.0 .+ 500.0 * randn(nx, ny))

            # Geopotential (increasing with height)
            g = 9.81
            for k in 1:6
                height = (k-1) * 1000.0  # Every 1km
                ds["PHB"][:, :, k, 1] .= Float32(g * height)
                ds["PH"][:, :, k, 1] = Float32.(g * 10.0 .+ g * 5.0 * randn(nx, ny))
            end
        end

        # Test reading mock file
        try
            bg_fields = WRFReader.read_wrf_netcdf(mock_file)

            @test bg_fields isa BackgroundFields
            @test size(bg_fields.u) == (8, 10, 5)  # (ny, nx, nz) after transpose
            @test size(bg_fields.v) == (8, 10, 5)
            @test size(bg_fields.t) == (8, 10, 5)
            @test size(bg_fields.q) == (8, 10, 5)
            @test size(bg_fields.ps) == (8, 10)
            @test size(bg_fields.z) == (8, 10, 5)

            # Check metadata
            @test bg_fields.metadata["model_type"] == "WRF-ARW"
            @test bg_fields.metadata["format"] == "NetCDF"
            @test bg_fields.metadata["grid_type"] == "lambert_conformal"
            @test bg_fields.metadata["map_proj"] == 1

            # Check time parsing
            @test bg_fields.time == DateTime(2018, 8, 12, 12, 0, 0)

            # Check physical range validity
            @test all(-100 .< bg_fields.u .< 100)  # m/s
            @test all(-100 .< bg_fields.v .< 100)  # m/s
            @test all(200 .< bg_fields.t .< 350)   # K
            @test all(0 .<= bg_fields.q .< 0.1)    # kg/kg
            @test all(50000 .< bg_fields.ps .< 105000)  # Pa
            @test all(bg_fields.z .>= 0)           # m

            println("✓ Mock WRF file test passed")
            println("  Grid dimensions: $(size(bg_fields.u))")
            println("  u range: $(extrema(bg_fields.u)) m/s")
            println("  v range: $(extrema(bg_fields.v)) m/s")
            println("  T range: $(extrema(bg_fields.t)) K")
            println("  q range: $(extrema(bg_fields.q)) kg/kg")
            println("  ps range: $(extrema(bg_fields.ps)) Pa")
            println("  z range: $(extrema(bg_fields.z)) m")

        finally
            # Clean up mock file
            rm(mock_file, force=true)
        end
    end

    @testset "Real WRF File Reading" begin
        if WRF_AVAILABLE
            println("\n📁 Real WRF file available at: $WRF_TEST_FILE")

            try
                bg_fields = WRFReader.read_wrf_netcdf(WRF_TEST_FILE)

                @test bg_fields isa BackgroundFields

                # Check that all fields have consistent dimensions
                nlat, nlon, nlev = size(bg_fields.u)
                @test size(bg_fields.v) == (nlat, nlon, nlev)
                @test size(bg_fields.t) == (nlat, nlon, nlev)
                @test size(bg_fields.q) == (nlat, nlon, nlev)
                @test size(bg_fields.ps) == (nlat, nlon)
                @test size(bg_fields.z) == (nlat, nlon, nlev)

                # Check metadata
                @test haskey(bg_fields.metadata, "model_type")
                @test haskey(bg_fields.metadata, "format")
                @test haskey(bg_fields.metadata, "grid_type")
                @test haskey(bg_fields.metadata, "source_file")

                # Check physical validity
                @test all(isfinite.(bg_fields.u))
                @test all(isfinite.(bg_fields.v))
                @test all(isfinite.(bg_fields.t))
                @test all(isfinite.(bg_fields.q))
                @test all(isfinite.(bg_fields.ps))
                @test all(isfinite.(bg_fields.z))

                # Temperature should be in reasonable range (K)
        @test all(190 .< bg_fields.t .< 350)

                # Specific humidity should be non-negative and < 1
                @test all(0 .<= bg_fields.q .< 1)

                # Surface pressure should be reasonable (Pa)
                @test all(30000 .< bg_fields.ps .< 110000)

                # Geopotential height should increase with level (roughly)
                @test bg_fields.z[1,1,end] > bg_fields.z[1,1,1]

                println("✓ Real WRF file test passed")
                println("  Source: $(bg_fields.metadata["source_file"])")
                println("  Grid dimensions: $nlat × $nlon × $nlev")
                println("  Valid time: $(bg_fields.time)")
                println("  Projection: $(bg_fields.metadata["grid_type"])")
                println("  Field ranges:")
                println("    u: $(extrema(bg_fields.u)) m/s")
                println("    v: $(extrema(bg_fields.v)) m/s")
                println("    T: $(extrema(bg_fields.t)) K")
                println("    q: $(extrema(bg_fields.q)) kg/kg")
                println("    ps: $(extrema(bg_fields.ps)) Pa")
                println("    z: $(extrema(bg_fields.z)) m")

            catch e
                @error "Failed to read real WRF file" exception=e
                rethrow(e)
            end
        else
            @warn "Real WRF test file not available: $WRF_TEST_FILE"
            println("⚠ Skipping real WRF file test - file not found")
        end
    end

    @testset "DataIO Integration" begin
        # Test that DataIO.read_background_fields works with WRF
        mock_file = joinpath(@__DIR__, "mock_wrf2.nc")

        # Create minimal mock file
        NCDataset(mock_file, "c") do ds
            defDim(ds, "Time", 1)
            defDim(ds, "west_east", 5)
            defDim(ds, "south_north", 4)
            defDim(ds, "bottom_top", 3)
            defDim(ds, "west_east_stag", 6)
            defDim(ds, "south_north_stag", 5)
            defDim(ds, "bottom_top_stag", 4)
            defDim(ds, "DateStrLen", 19)

            defVar(ds, "XLAT", Float32, ("west_east", "south_north", "Time"))
            defVar(ds, "XLONG", Float32, ("west_east", "south_north", "Time"))
            defVar(ds, "Times", Char, ("DateStrLen", "Time"))
            defVar(ds, "U", Float32, ("west_east_stag", "south_north", "bottom_top", "Time"))
            defVar(ds, "V", Float32, ("west_east", "south_north_stag", "bottom_top", "Time"))
            defVar(ds, "T", Float32, ("west_east", "south_north", "bottom_top", "Time"))
            defVar(ds, "QVAPOR", Float32, ("west_east", "south_north", "bottom_top", "Time"))
            defVar(ds, "P", Float32, ("west_east", "south_north", "bottom_top", "Time"))
            defVar(ds, "PB", Float32, ("west_east", "south_north", "bottom_top", "Time"))
            defVar(ds, "PSFC", Float32, ("west_east", "south_north", "Time"))
            defVar(ds, "PH", Float32, ("west_east", "south_north", "bottom_top_stag", "Time"))
            defVar(ds, "PHB", Float32, ("west_east", "south_north", "bottom_top_stag", "Time"))

            ds.attrib["MAP_PROJ"] = 1
            ds.attrib["DX"] = 9000.0
            ds.attrib["DY"] = 9000.0
            ds.attrib["CEN_LAT"] = 40.0
            ds.attrib["CEN_LON"] = -95.0

            # Fill with data
            ds["U"][:,:,:,1] = randn(Float32, 6, 4, 3)
            ds["V"][:,:,:,1] = randn(Float32, 5, 5, 3)
            ds["T"][:,:,:,1] = Float32.(5.0 .+ randn(5, 4, 3))
            ds["QVAPOR"][:,:,:,1] = Float32.(abs.(0.005 .+ 0.001*randn(5, 4, 3)))
            ds["P"][:,:,:,1] = randn(Float32, 5, 4, 3) .* 100
            ds["PB"][:,:,:,1] .= Float32(100000.0)
            ds["PSFC"][:,:,1] = Float32.(101325.0 .+ randn(5, 4))
            ds["PH"][:,:,:,1] = randn(Float32, 5, 4, 4) .* 100
            ds["PHB"][:,:,:,1] = Float32.([i*1000*9.81 for i in 0:3, j in 1:4, k in 1:5])
            ds["XLAT"][:,:,1] = Float32.(repeat(30:33, 1, 5)')
            ds["XLONG"][:,:,1] = Float32.(repeat(-100:-96, 4, 1))
            ds["Times"][:,1] = collect("2018-08-12_12:00:00")
        end

        try
            # Test integration through DataIO interface
            bg_fields = DataIO.read_background_fields(mock_file, format=:netcdf, model_type=:wrf)

            @test bg_fields isa BackgroundFields
            @test size(bg_fields.u) == (4, 5, 3)

            println("✓ DataIO integration test passed")

        finally
            rm(mock_file, force=true)
        end
    end
end

println("\n" * "="^60)
println("WRF Reader Test Summary")
println("="^60)
if WRF_AVAILABLE
    println("✓ All tests passed (including real WRF file)")
else
    println("✓ All mock tests passed")
    println("⚠ Real WRF file tests skipped (file not available)")
end
println("="^60)
